import config
import tensorflow as tf
import tensorflow_addons as tfa


class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.embed = tf.keras.layers.Embedding(
            config.VOCAB_SIZE,
            config.EMBEDDING_DIM,
        )

        self.lstm = tf.keras.layers.Bidirectional(
            tf.keras.layers.LSTM(config.HIDDEN_SIZE, return_sequences=True)
        )
        self.linear = tf.keras.layers.Dense(config.TARGET_SIZE)

        # CRF layer
        self.crf = tfa.layers.CRF(config.TARGET_SIZE)

    def _get_bi_lstm_feature(self, x):
        out = self.embed(x)
        out, _ = self.lstm(out)
        return self.linear(out)

    def call(self, inputs, training=None, mask=None):
        out = self._get_bi_lstm_feature(inputs)
        return self.crf.decode(out, mask)

    def loss_fn(self, x, target, mask):
        y_pred = self._get_bi_lstm_feature(x)
        return -self.crf(y_pred, target, mask, reduction='mean')
